#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Qwen3-0.6B 简单思考模式演示
基于您提供的例子，实现最简化的思考过程显示
"""

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

def main():
    """主函数 - 基于您提供的例子"""
    print("=== Qwen3-0.6B 思考模式演示 ===")
    
    # 设置模型本地路径
    model_name = "./modelscope_cache/Qwen/Qwen3-0___6B"
    
    print(f"正在加载模型: {model_name}")
    
    try:
        # 加载分词器和模型
        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype="auto",  # 自动选择合适的数据类型
            device_map="cuda:0" if torch.cuda.is_available() else "cpu",    # 自动选择可用设备(CPU/GPU)
            trust_remote_code=True
        )
        
        print("模型加载完成！")
        
        # 准备模型输入
        prompt = "你好，请介绍一下自己"
        messages = [
            {"role": "user", "content": prompt}
        ]
        
        print(f"用户问题: {prompt}")
        print("=" * 50)
        
        text = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
            enable_thinking=True # 选择是否打开深度推理模式
        )
        
        # 将输入文本转换为模型可处理的张量格式
        model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
        
        print("正在生成回复...")
        
        # 生成文本
        generated_ids = model.generate(
            **model_inputs,
            max_new_tokens=32768  # 设置最大生成token数量
        )
        
        # 提取新生成的token ID
        output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist() 
        
        # 解析思考内容
        try:
            # rindex finding 151668 (</think>)
            # 查找结束标记"</think>"的位置
            index = len(output_ids) - output_ids[::-1].index(151668)
        except ValueError:
            index = 0
        
        # 解码思考内容和最终回答
        thinking_content = tokenizer.decode(output_ids[:index], skip_special_tokens=True).strip("\n")
        content = tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n")
        
        # 打印结果
        if thinking_content:
            print("🤔 思考过程:")
            print(thinking_content)
            print("\n" + "=" * 50)
        
        print("💬 最终回答:")
        print(content)
        
    except FileNotFoundError:
        print("❌ 模型文件未找到！")
        print("请确保模型已下载到正确路径，或运行完整的下载脚本")
        print("模型路径:", model_name)
        
    except Exception as e:
        print(f"❌ 发生错误: {e}")
        print("\n🔧 可能的解决方案:")
        print("1. 检查模型路径是否正确")
        print("2. 确保已安装必要的依赖包")
        print("3. 检查CUDA是否可用（如果使用GPU）")

if __name__ == "__main__":
    main()
